import cv2
import subprocess
import time

# RTSP流地址和RTMP服务器地址
RTSP_URL = "rtsp://telchina:123qweASD@10.10.30.39:554//Streaming/Channels/1"
RTMP_URL = "rtmp://10.10.50.127:10935/live/stream_5"

def stream_video():
    while True:  # 持续尝试推送视频流
        try:
            # 创建VideoCapture对象
            cap = cv2.VideoCapture(RTSP_URL)

            # 检查视频流是否成功打开
            if not cap.isOpened():
                print("无法打开RTSP流，等待后重试...")
                time.sleep(5)  # 等待5秒后重试
                continue

            # 获取视频的宽度和高度
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

            # 设置ffmpeg推流参数
            command = [
                'ffmpeg',
                '-y',  # 覆盖输出文件
                '-f', 'rawvideo',
                '-pix_fmt', 'bgr24',
                '-s', f'{width}x{height}',  # 输出视频分辨率
                '-r', '30',  # 帧率
                '-i', '-',  # 从stdin读取视频流
                '-c:v', 'libx264',  # 使用x264编码
                '-preset', 'veryfast',  # 编码速度
                '-tune', 'zerolatency',  # 针对低延迟的调优
                '-b:v', '1000k',  # 设置视频比特率
                '-maxrate', '1000k',  # 最大比特率
                '-bufsize', '2000k',  # 缓冲区大小
                '-f', 'flv',  # 输出格式
                RTMP_URL
            ]

            # 启动ffmpeg进程
            process = subprocess.Popen(command, stdin=subprocess.PIPE)

            while True:
                ret, frame = cap.read()
                
                if not ret:
                    print("读取帧失败，结束流处理")
                    break

                # 绘制识别框（示例为固定位置的框）
                cv2.rectangle(frame, (50, 50), (200, 200), (0, 255, 0), 2)

                # 显示当前帧（可选）
                cv2.imshow("Video Stream", frame)

                # 将帧写入ffmpeg进程的stdin
                process.stdin.write(frame.tobytes())

                # 按 'q' 键退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

            # 处理结束后，确保关闭ffmpeg进程
            process.stdin.close()
            process.wait()

        except Exception as e:
            print(f"发生异常: {e}")
        
        finally:
            cap.release()
            cv2.destroyAllWindows()
            print("尝试重新推送视频流...")
            time.sleep(5)  # 等待5秒后重试

if __name__ == "__main__":
    stream_video()
